import matplotlib as mpl
mpl.use('Agg')
import sdf_helper as sh
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib as matplb
import numpy as np
import argparse
import os
parser = argparse.ArgumentParser()
parser.add_argument('--file', type=str, default=None)
parser.add_argument('--method', type=str, default='scatter')
pargs = parser.parse_args()
if (os.path.exists(pargs.file)):
    data = sh.getdata(pargs.file)
else:
    print('cannot find the file')
    exit(1)
var = data.Electric_Field_Ey
xmax = len(var.grid_mid.data[0])
ymax = len(var.grid_mid.data[1])
zmax = len(var.grid_mid.data[2])
xslice = range(0, int(xmax / 10))
yslice = range(0, int(ymax / 10))
zslice = range(0, int(zmax / 10))
x = np.array([])
y = np.array([])
z = np.array([])
c = np.array([])
for i in xslice:
    for j in yslice:
        for k in zslice:
            x = np.append(x, var.grid_mid.data[0][i])
            y = np.append(y, var.grid_mid.data[1][j])
            z = np.append(z, var.grid_mid.data[2][k])
            c = np.append(c, var.data[i][j][k])
print(len(x))
print(x)
print(len(y))
print(len(z))
print(c)
fig = plt.figure(figsize=(15, 10))
ax = fig.add_subplot(111, projection='3d')
print('ploting')
cax = getattr(ax, pargs.method)(x, y, z, c=c, cmap='jet')
plt.colorbar(cax, ax=ax)
plt.savefig('3d-profile.png')